import concurrent.futures
import datetime
import functools
import logging
import traceback
import pandas
import tabulate

from absl import (app, flags)

from coin.base.datetime_util import (
    to_timestamp_int, convert_string_to_dates, iterate_date)
from coin.base.param_util import to_list
from coin.support.proto_log.logic.util import read_strat_log
from coin.support.pta.app.pta_dumper_runner import (
    _get_dt_ranges, _get_strat_info_list, _parse_time)
from coin.tool.strat_monitor.app.alarm_strategy_log import StrategyLogAlarm
from coin.tool.strat_monitor.util.monitor_util import (get_active_strategy,
                                                       get_strat_from_json_config)
from coin.util.queue.config import KafkaConfig
from xunkemgmt_client.client.util.query_util import query_strategies


FLAGS = flags.FLAGS

WARMUP_OFFSET_TS = 60 * 60 * 10**9


def _dump_alarms(alarms, output_file):
  df = pandas.DataFrame(alarms)
  alarm_str = tabulate.tabulate(
      df.sort_values(['alarm_type', 'strategy_name']).reset_index(drop=True),
      headers="keys")
  if output_file is not None:
    with open(output_file, "w") as fp:
      fp.write(alarm_str)
  else:
    print(alarm_str)


class IntervalContainer(object):
  def __init__(self, start_ts, end_ts, frequency):
    assert start_ts < end_ts, (start_ts, end_ts)
    intervals = []
    ts = start_ts
    intervals.append(start_ts)
    while ts < end_ts:
      ts += frequency
      if ts <= end_ts:
        intervals.append(ts)
      else:
        intervals.append(end_ts)
    self._intervals = intervals

  def get_interval(self, value):
    size = len(self._intervals)
    for i in range(size):
      if i != size - 1:
        if self._intervals[i] <= value < self._intervals[i+1]:
          interval = (self._intervals[i], self._intervals[i+1])
          break
    else:
      raise ValueError(f'interval not found: {value}')
    return interval

  def get_enclosing_interval(self, start, end):
    assert start < end, (start, end)
    start_idx = self._intervals.index(self.get_interval(start)[1])
    end_idx = self._intervals.index(self.get_interval(end)[1])
    intervals = []
    while start_idx < end_idx:
      intervals.append(self._intervals[start_idx])
      start_idx += 1
    return intervals


class StrategyAlarmReplayer(object):
  def __init__(self, *, strat_info, start_ts, end_ts, alarm_frequency, reset_frequency):
    self._strat_info = strat_info
    self._alarm_interval = IntervalContainer(
        start_ts, end_ts, int(alarm_frequency * 10**9))
    self._reset_interval = IntervalContainer(
        start_ts, end_ts, int(reset_frequency * 10**9))
    self._current_ts = start_ts
    self._prev_reset_ts = self._alarm_interval.get_interval(start_ts)[0]
    self._next_alarm_ts = self._reset_interval.get_interval(start_ts)[1]
    self._alarm_monitor = StrategyLogAlarm(
        strat_infos=to_list(strat_info),
        current_ts=start_ts)
    self._alarms = []

  @property
  def alarms(self):
    return self._alarms

  def _gen_alarm_output(self, alarm_infos, alarm_ts):
    alarms = []
    for alarm_type, alarm_info in alarm_infos.items():
      for strat_name, alarm_msg in alarm_info.items():
        alarms.append({
            'strategy_name': strat_name,
            'timestamp': alarm_ts,
            'time': datetime.datetime.utcfromtimestamp(
                alarm_ts / 10**9).strftime('%Y%m%d %H%M%S'),
            'alarm_type': alarm_type.__name__,
            'alarm_msg': alarm_msg,
        })
    return alarms

  def on_log(self, timestamp, log):
    if timestamp < self._current_ts:
      return
    if timestamp > self._next_alarm_ts:
      self._next_alarm_ts = self._alarm_interval.get_interval(timestamp)[1]
      alarm_intervals = self._alarm_interval.get_enclosing_interval(
          self._current_ts, timestamp)
      for alarm_ts in alarm_intervals:
        assert alarm_ts >= self._current_ts, (alarm_ts, self._current_ts)
        self._current_ts = alarm_ts
        prev_reset_ts = self._reset_interval.get_interval(alarm_ts)[0]
        if prev_reset_ts != self._prev_reset_ts:
          self._alarm_monitor._reset()
          self._prev_reset_ts = prev_reset_ts
        self._alarm_monitor.update_all_timestamp(alarm_ts)
        alarm_infos = self._alarm_monitor.check_alarms()
        self._alarms += self._gen_alarm_output(alarm_infos, alarm_ts)
    self._alarm_monitor.on_log(
        self._strat_info.strategy_name,
        self._strat_info.machine,
        timestamp,
        log)
    self._current_ts = timestamp

  def on_kafka_log(self, partition, record):
    self.on_log(record.timestamp * 1e6, record.value)


class AlarmDumperRunner(object):
  def __init__(self,
               *,
               strat_list,
               kafka_config_filename,
               strat_log_root_dir,
               max_workers,
               **alarm_params):
    self._start_ts = alarm_params.get('start_ts')
    self._end_ts = alarm_params.get('end_ts')
    self._strat_log_root_dir = strat_log_root_dir
    self._max_workers = max_workers
    self._alarms = []
    self._strat_log_root_dir = strat_log_root_dir
    self._alarm_params = alarm_params
    self._kafka_config = None
    if kafka_config_filename is not None:
      self._kafka_config = KafkaConfig.from_cmd_config(kafka_config_filename)
    start_date = datetime.datetime.utcfromtimestamp(self._start_ts / 10**9).date()
    end_date = datetime.datetime.utcfromtimestamp(self._end_ts / 10**9).date()
    self._strat_infos = _get_strat_info_list(
        start_date=start_date, end_date=end_date, strategy_list=strat_list,
        kafka_config=self._kafka_config, strat_log_root_dir=strat_log_root_dir)

  def _replay_alarm(self, *, strat_info):
    alarm_replayer = StrategyAlarmReplayer(**self._alarm_params, strat_info=strat_info)
    if self._kafka_config is not None:
      callback = alarm_replayer.on_kafka_log
    else:
      callback = alarm_replayer.on_log
    start_ts = self._start_ts - WARMUP_OFFSET_TS
    read_strat_log(
        strat_info=strat_info,
        start_time=datetime.datetime.utcfromtimestamp(start_ts / 10**9),
        end_time=datetime.datetime.utcfromtimestamp(self._end_ts / 10**9),
        callback=callback,
        kafka_config=self._kafka_config,
        root_dir=self._strat_log_root_dir)
    return alarm_replayer.alarms

  @property
  def alarms(self):
    return self._alarms

  def reset_alarms(self):
    self._alarms = []

  def run(self):
    func_list = []
    for strat_info in self._strat_infos:
      replay_func = functools.partial(self._replay_alarm, strat_info=strat_info)
      func_list.append((strat_info, replay_func))
    num_func = len(func_list)
    self.reset_alarms()
    alarms = []
    success = 0
    max_workers = self._max_workers
    if max_workers > 1:
      with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
        fut_to_strat_info = {executor.submit(func): key for key, func in func_list}
        done_futs, not_done_futs = concurrent.futures.wait(fut_to_strat_info.keys(),
                                                           timeout=None)
        func_list = [(fut_to_strat_info[fut], fut.result) for fut in done_futs]
        for fut in not_done_futs:
          strat_info = fut_to_strat_info[fut]
          logging.error('Fail to dump strategy log due to timeout. %s',
                        strat_info._asdict())
        for pid, process in executor._processes.items():
          process.terminate()

    for key, func in func_list:
      try:
        alarms.extend(func())
        success += 1
      except Exception as e:
        logging.error('Fail to dump strategy log. %s, %s, %s', key._asdict(), type(e), e)
        logging.error(traceback.format_exc())
    logging.info('Success rate: %s/%s' % (success, num_func))
    num_fail = num_func - success
    self._alarms = alarms
    return num_fail


def main(_):
  start_time = FLAGS.start_time
  assert start_time, '--start_time must be specified.'
  end_time = FLAGS.end_time
  assert end_time, '--end_time must be specified.'
  start_dt = _parse_time(start_time)
  end_dt = _parse_time(end_time)
  dt_ranges = _get_dt_ranges(start_dt, end_dt, FLAGS.aggregate)

  strategy_universe = query_strategies(as_proto=True)
  strategy_config_filename = FLAGS.strategy_config_filename
  if strategy_config_filename is not None:
    strat_list = get_strat_from_json_config(strategy_config_filename)
  else:
    threshold = to_timestamp_int(start_dt)
    strat_list = get_active_strategy(update_timestamp_threshold=threshold)
  filtered = strategy_universe
  if FLAGS.business_unit is not None:
    business_units = [elem.strip() for elem in FLAGS.business_unit.split(',')]
    filtered = [elem for elem in filtered if elem.business_unit in business_units]
  if FLAGS.strategy_group is not None:
    strategy_groups = [elem.strip() for elem in FLAGS.strategy_group.split(',')]
    filtered = [elem for elem in filtered if elem.strategy_group in strategy_groups]
  strat_list = set(strat_list).intersection(
      set([elem.strategy_name for elem in filtered]))

  for start_dt, end_dt in dt_ranges:
    start_td_str = start_dt.strftime('%Y%m%dT%H%M%S')
    end_td_str = end_dt.strftime('%Y%m%dT%H%M%S')
    logging.info('\nRunning for %s-%s ...' % (start_td_str, end_td_str))
    start_ts = to_timestamp_int(start_dt)
    end_ts = to_timestamp_int(end_dt)
    alarm_dumper_runner = AlarmDumperRunner(
        strat_list=strat_list,
        kafka_config_filename=FLAGS.kafka_config_filename,
        strat_log_root_dir=FLAGS.root_dir,
        max_workers=FLAGS.max_workers,
        start_ts=start_ts,
        end_ts=end_ts,
        alarm_frequency=FLAGS.alarm_frequency,
        reset_frequency=FLAGS.reset_frequency)
    alarm_dumper_runner.run()
    _dump_alarms(alarm_dumper_runner.alarms, output_file=FLAGS.output_file)


if __name__ == '__main__':
  logging.basicConfig(level='INFO', format='%(levelname)8s %(asctime)s %(message)s')

  flags.DEFINE_string('start_time', None, 'Start date in form of %Y%m%d.')
  flags.DEFINE_string('end_time', None, 'End date in form of %Y%m%d.')
  flags.DEFINE_integer('max_workers', 24, 'Max number of workers.')
  flags.DEFINE_string('root_dir', None, '')
  flags.DEFINE_boolean('aggregate', False, '')
  flags.DEFINE_string('kafka_config_filename', None, 'kafka config')
  flags.DEFINE_string('strategy_config_filename', None, 'strategy_config_filename')
  flags.DEFINE_string('strategy_group', None, 'filter by strategy_group')
  flags.DEFINE_string('business_unit', None, 'filter by business_unit')
  flags.DEFINE_string('output_file', None, 'output_file')
  flags.DEFINE_integer('alarm_frequency', 1800, 'alarm frequency')
  flags.DEFINE_integer('reset_frequency', 1800, 'alarm frequency')
  app.run(main)
